import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from MnistConvNet import ConvNet

# Device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Hyper parameters
num_epochs = 5
num_classes = 10
batch_size = 100
learning_rate = 0.001

# MNIST dataset
'''
train（bool，可选）–如果为True，则从training.pt创建数据集，否则从test.pt创建数据集
transforms.ToTensor()将numpy的ndarray或PIL.Image读的图片转换成形状为(C,H, W)的Tensor格式，
且/255归一化到[0,1.0]之间
'''
train_dataset = torchvision.datasets.MNIST(root='./resources/',
                                           train=True,
                                           transform=transforms.ToTensor(),
                                           download=False)

test_dataset = torchvision.datasets.MNIST(root='./resources/',
                                          train=False,
                                          transform=transforms.ToTensor())

# Data loader
# 从已有的数据读取接口的输入按照batch size封装成Tensor,shuffle=True为每迭代一次洗牌重新按照batch_size分
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)

model = ConvNet(num_classes).to(device)

# Loss and optimizer，交叉熵损失函数，Adam是比较好的优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# Train the model,总的批数
total_step = len(train_loader)

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward and optimize，将梯度置零每一批次后，避免梯度叠加,因为使用的是当前参数空间对应的梯度
        optimizer.zero_grad()
        # 损失后向传播，计算梯度
        loss.backward()
        # 优化器更新参数
        optimizer.step()

        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                  .format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))

# Test the model
'''
model.eval()是保证BN层能够用全部训练数据的均值和方差，即测试过程中要保证BN层的均值和方差不变。
对于Dropout，model.eval()是利用到了所有网络连接，即不进行随机舍弃神经元。
'''

model.eval()  # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
# 强制之后的内容不进行计算图构建
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        # 第一个值是具体的value（我们用下划线_表示），第二个值是value所在的index（也就是predicted）,1为所在行的最大值
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))

# Save the model checkpoint 保存训练好的参数
torch.save(model.state_dict(), 'ckpt/MnistModel.ckpt')

